#!/usr/bin/env python2
from __future__ import division

from pylab import *
from numpy import * #unnecessary, just for pydev auto-complete
from matching_pursuit import *
from fJ_generator import *
from scipy import optimize
from scikits.audiolab import wavread
from scikits.audiolab import wavwrite
from scikits.audiolab import play
from levenberg_marquardt import ooopti
#from enthought.mayavi.mlab import *
import os

def rotate(l, n):
    return l[n:] + l[:n]
framesize = 44100 / 50


class _CLASS_:
    def __init__(self):
        self.t = sympy.Symbol('t')
        self.p = sympy.Symbol('p')
        self.A_l = sympy.Symbol('A_l')
        self.A_r = sympy.Symbol('A_r')
        self.w_l = sympy.Symbol('w_l')
        self.w_r = sympy.Symbol('w_r')

        self.fJ_linear_full = fJ_generator(sympy.exp(self.A_l * (1 - self.t) + self.A_r * self.t) *
                                           sympy.sin(self.p + self.t * sympy.exp(self.w_l * (1 - self.t) + self.w_r * self.t)),
                                           [self.A_l, self.A_r, self.w_l, self.w_r, self.p],
                                           {self.t:linspace(0, 1, framesize, endpoint=False)})
    def optimize_linear_full(self, x0, data=None):
        (result, success, r1, r2, r3) = optimize.leastsq(self.fJ_linear_full.f, x0, args=(data), Dfun=self.fJ_linear_full.J, full_output=True)
        return result
#    def optimize_linear_l(self, x0, data=None):
#        self.fJ_linear_l = fJ_generator(sympy.exp(self.A_l * (1 - self.t) + self.A_r * self.t) *
#                                        sympy.sin(self.p + self.t * sympy.exp(self.w_l * (1 - self.t) + self.w_r * self.t)),
#                                        [self.A_l, self.A_r, self.w_l, self.w_r, self.p],
#                                        {self.t:linspace(0, 1, framesize, endpoint=False)})
#        (result, success, r1, r2, r3) = optimize.leastsq(self.fJ_linear_l.f, x0, args=(data), Dfun=self.fJ_linear_l.J, full_output=True)
        return result
    def optimize_linear_r(self, x0, data=None):
        [A_l, A_r, w_l, w_r, p] = x0
        self.fJ_linear_r = fJ_generator(sympy.exp(A_l * (1 - self.t) + self.A_r * self.t) *
                                        sympy.sin(self.p + self.t * sympy.exp(w_l * (1 - self.t) + self.w_r * self.t)),
                                        [self.A_r, self.w_r, self.p],
                                        {self.t:linspace(0, 1, framesize, endpoint=False)})
        (result, success, r1, r2, r3) = optimize.leastsq(self.fJ_linear_r.f, [A_r, w_r, p], args=(data), Dfun=self.fJ_linear_r.J, full_output=True)
        result = asarray([A_l, result[0], w_l, result[1], result[2]])
        return result

class DATA:
    def __init__(self, file):
        wav, fs, enc = wavread(file)
        if len(wav.shape) == 1:
            self.data = wav[:]
        elif len(wav.shape) == 2:
            self.data = wav.T[0, :]
    def getframe(self, location):
        return self.data[location:location + framesize]

if __name__ == '__main__':
    testsound = [None]*250
    _X_ = _CLASS_()
    _DATA_ = DATA("/home/kkb110/music/research/samples/VI_sV_espr_4s_ff_D5.wav")
    startingpoint = 0 * 44100
    data = _DATA_.getframe(startingpoint)

    mp = MatchingPursuit(framesize, 12.5, 120 * 10, 120)
    mp_A, mp_w, mp_p = mp.matching_pursuit(data)
    mp_p += pi / 2


    p = mp_p
    w_l = log(mp_w)
    w_r = log(mp_w)
    A_l = log(mp_A)
    A_r = log(mp_A)

    [A_l, A_r, w_l, w_r, p] = _X_.optimize_linear_full([A_l, A_r, w_l, w_r, p], data)

    plot(data)
    plot(_X_.fJ_linear_full.f([A_l, A_r, w_l, w_r, p]))
    testsound[0] = _X_.fJ_linear_full.f([A_l, A_r, w_l, w_r, p])
    show()

    for i in range(1, 250):
        data = _DATA_.getframe(startingpoint + i * framesize)
        [A_l, A_r, w_l, w_r, p] = _X_.optimize_linear_r([A_r, A_r, w_r, w_r, p], data)
        plot(data)
        plot(_X_.fJ_linear_full.f([A_l, A_r, w_l, w_r, p]))
        testsound[i] = _X_.fJ_linear_full.f([A_l, A_r, w_l, w_r, p])
        #show()
    #play(hstack(testsound))
    wavwrite(hstack(testsound), "/home/kkb110/music/research/samples/VI_sV_espr_4s_ff_D5_test1.wav", fs=44100)














